// COLMAP - Structure-from-Motion and Multi-View Stereo.
// Copyright (C) 2016  Johannes L. Schoenberger <jsch at inf.ethz.ch>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#ifndef COLMAP_SRC_OPTIM_LORANSAC_H_
#define COLMAP_SRC_OPTIM_LORANSAC_H_

#include <cfloat>
#include <random>
#include <stdexcept>
#include <vector>

#include <Eigen/Core>

#include "optim/random_sampler.h"
#include "optim/ransac.h"
#include "optim/support_measurement.h"

namespace RVG
{

// Implementation of LO-RANSAC (Locally Optimized RANSAC).
//
// "Locally Optimized RANSAC" Ondrej Chum, Jiri Matas, Josef Kittler, DAGM 2003.
template <typename Estimator, typename LocalEstimator,
         typename SupportMeasurer = InlierSupportMeasurer,
         typename Sampler = RandomSampler>
class LORANSAC : public RANSAC<Estimator, SupportMeasurer, Sampler>
{
public:
    using typename RANSAC<Estimator, SupportMeasurer, Sampler>::Report;

    LORANSAC(const RANSACOptions& options);

    // Robustly estimate model with RANSAC (RANdom SAmple Consensus).
    //
    // @param X              Independent variables.
    // @param Y              Dependent variables.
    //
    // @return               The report with the results of the estimation.
    Report Estimate(const std::vector<typename Estimator::X_t>& X,
                    const std::vector<typename Estimator::Y_t>& Y);

    Report Estimate_P2P(const std::vector<typename Estimator::X_t>& X,
                        const std::vector<typename Estimator::Y_t>& Y,
                        const Eigen::Matrix3d& Rprior);
    // Objects used in RANSAC procedure.
    using RANSAC<Estimator, SupportMeasurer, Sampler>::estimator;
    LocalEstimator local_estimator;
    using RANSAC<Estimator, SupportMeasurer, Sampler>::sampler;
    using RANSAC<Estimator, SupportMeasurer, Sampler>::support_measurer;

private:
    using RANSAC<Estimator, SupportMeasurer, Sampler>::options_;
};

////////////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////////////

template <typename Estimator, typename LocalEstimator, typename SupportMeasurer,
         typename Sampler>
LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::LORANSAC(
    const RANSACOptions& options)
    : RANSAC<Estimator, SupportMeasurer, Sampler>(options) {}

template <typename Estimator, typename LocalEstimator, typename SupportMeasurer,
         typename Sampler>
typename LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::Report
LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::Estimate_P2P(
    const std::vector<typename Estimator::X_t>& X,
    const std::vector<typename Estimator::Y_t>& Y,
    const Eigen::Matrix3d& Rprior)
{
    CHECK_EQ(X.size(), Y.size());

    const size_t num_samples = X.size();

    typename RANSAC<Estimator, SupportMeasurer, Sampler>::Report report;
    report.success = false;
    report.num_trials = 0;

    if (num_samples < Estimator::MinNumSamples())
    {
        return report;
    }

    typename SupportMeasurer::Support best_support;
    typename Estimator::M_t best_model;
    bool best_model_is_local = false;

    bool abort = false;

    const double max_residual = options_.max_error * options_.max_error;

    std::vector<double> residuals(num_samples);

    std::vector<typename LocalEstimator::X_t> X_inlier;
    std::vector<typename LocalEstimator::Y_t> Y_inlier;

    std::vector<typename Estimator::X_t> X_rand(Estimator::MinNumSamples());
    std::vector<typename Estimator::Y_t> Y_rand(Estimator::MinNumSamples());

    sampler.Initialize(num_samples);
    size_t max_num_trials = options_.max_num_trials;
    max_num_trials = std::min<size_t>(max_num_trials, sampler.MaxNumSamples());
    size_t dyn_max_num_trials = max_num_trials;

    std::vector<typename Estimator::Y_t> RY;
    for(size_t iY = 0; iY < Y.size(); iY++)
    {
        RY.push_back(Rprior * Y[iY]);
    }

    //LOG(INFO) << "R transformed " << std::endl;
    for (report.num_trials = 0; report.num_trials < max_num_trials;
            ++report.num_trials)
    {
        if (abort)
        {
            report.num_trials += 1;
            break;
        }

        sampler.SampleXY(X, RY, &X_rand, &Y_rand); // X is new R^T * inv(K) * p

        // Estimate model for current subset.
        std::vector<typename Estimator::M_t> sample_models = estimator.Estimate(X_rand, Y_rand);

        for (auto& sample_model : sample_models)
        {
            //LOG(INFO) << sample_model.col(3)(0) << " " << sample_model.col(3)(1) << " "<< sample_model.col(3)(2) << std::endl;
            sample_model(0, 0) = Rprior(0, 0); // update the prior rotation matrix 2017-05-01
            sample_model(0, 1) = Rprior(0, 1);
            sample_model(0, 2) = Rprior(0, 2);
            sample_model(1, 0) = Rprior(1, 0);
            sample_model(1, 1) = Rprior(1, 1);
            sample_model(1, 2) = Rprior(1, 2);
            sample_model(2, 0) = Rprior(2, 0);
            sample_model(2, 1) = Rprior(2, 1);
            sample_model(2, 2) = Rprior(2, 2);
        }

        // Iterate through all estimated models
        for (const auto& sample_model : sample_models)
        {
            estimator.Residuals(X, Y, sample_model, &residuals);
            CHECK_EQ(residuals.size(), X.size());

            const auto support = support_measurer.Evaluate(residuals, max_residual);

            // Do local optimization if better than all previous subsets.
            if (support_measurer.Compare(support, best_support))
            {
                best_support = support;
                best_model = sample_model;
                best_model_is_local = false;

                // Estimate locally optimized model from inliers.
                /*if (support.num_inliers > Estimator::MinNumSamples() &&
                    support.num_inliers >= LocalEstimator::MinNumSamples()) {
                  X_inlier.clear();
                  Y_inlier.clear();
                  X_inlier.reserve(support.num_inliers);
                  Y_inlier.reserve(support.num_inliers);
                  for (size_t i = 0; i < residuals.size(); ++i) {
                    if (residuals[i] <= max_residual) {
                      X_inlier.push_back(X[i]);
                      Y_inlier.push_back(Y[i]);
                    }
                  }

                  //const std::vector<typename LocalEstimator::M_t> local_models =
                  //    local_estimator.Estimate(X_inlier, Y_inlier);

                  //for (const auto& local_model : local_models) {
                  //  local_estimator.Residuals(X, Y, local_model, &residuals);
                  //  CHECK_EQ(residuals.size(), X.size());

                 //   const auto local_support =
                 //       support_measurer.Evaluate(residuals, max_residual);

                    // Check if non-locally optimized model is better.
                 //   if (support_measurer.Compare(local_support, support)) {
                 //     best_support = local_support;
                 //     best_model = local_model;
                 //     best_model_is_local = true;
                 //   }
                  }
                }*/

                if (report.num_trials >= options_.min_num_trials)
                {
                    dyn_max_num_trials =
                        RANSAC<Estimator, SupportMeasurer, Sampler>::ComputeNumTrials(
                            best_support.num_inliers, num_samples, options_.confidence);
                }
            }

            if (report.num_trials >= dyn_max_num_trials)
            {
                abort = true;
                break;
            }
        }
    }

    report.support = best_support;
    report.model = best_model;

    // No valid model was found
    if (report.support.num_inliers < estimator.MinNumSamples())
    {
        return report;
    }

    report.success = true;

    // Determine inlier mask. Note that this calculates the residuals for the
    // best model twice, but saves to copy and fill the inlier mask for each
    // evaluated model. Some benchmarking revealed that this approach is faster.

    if (best_model_is_local)
    {
        local_estimator.Residuals(X, Y, report.model, &residuals);
    }
    else
    {
        estimator.Residuals(X, Y, report.model, &residuals);
    }

    CHECK_EQ(residuals.size(), X.size());

    report.inlier_mask.resize(num_samples);
    for (size_t i = 0; i < residuals.size(); ++i)
    {
        if (residuals[i] <= max_residual)
        {
            report.inlier_mask[i] = true;
        }
        else
        {
            report.inlier_mask[i] = false;
        }
    }

    return report;
}



template <typename Estimator, typename LocalEstimator, typename SupportMeasurer,
         typename Sampler>
typename LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::Report
LORANSAC<Estimator, LocalEstimator, SupportMeasurer, Sampler>::Estimate(
    const std::vector<typename Estimator::X_t>& X,
    const std::vector<typename Estimator::Y_t>& Y)
{
    CHECK_EQ(X.size(), Y.size());

    const size_t num_samples = X.size();

    typename RANSAC<Estimator, SupportMeasurer, Sampler>::Report report;
    report.success = false;
    report.num_trials = 0;

    if (num_samples < Estimator::MinNumSamples())
    {
        return report;
    }

    typename SupportMeasurer::Support best_support;
    typename Estimator::M_t best_model;
    bool best_model_is_local = false;

    bool abort = false;

    const double max_residual = options_.max_error * options_.max_error;

    std::vector<double> residuals(num_samples);

    std::vector<typename LocalEstimator::X_t> X_inlier;
    std::vector<typename LocalEstimator::Y_t> Y_inlier;

    std::vector<typename Estimator::X_t> X_rand(Estimator::MinNumSamples());
    std::vector<typename Estimator::Y_t> Y_rand(Estimator::MinNumSamples());

    sampler.Initialize(num_samples);

    size_t max_num_trials = options_.max_num_trials;
    max_num_trials = std::min<size_t>(max_num_trials, sampler.MaxNumSamples());
    size_t dyn_max_num_trials = max_num_trials;

    for (report.num_trials = 0; report.num_trials < max_num_trials;
            ++report.num_trials)
    {
        if (abort)
        {
            report.num_trials += 1;
            break;
        }

        sampler.SampleXY(X, Y, &X_rand, &Y_rand);

        // Estimate model for current subset.
        const std::vector<typename Estimator::M_t> sample_models = estimator.Estimate(X_rand, Y_rand);

        // Iterate through all estimated models
        for (const auto& sample_model : sample_models)
        {
            estimator.Residuals(X, Y, sample_model, &residuals);
            CHECK_EQ(residuals.size(), X.size());

            const auto support = support_measurer.Evaluate(residuals, max_residual);

            // Do local optimization if better than all previous subsets.
            if (support_measurer.Compare(support, best_support))
            {
                best_support = support;
                best_model = sample_model;
                best_model_is_local = false;

                // Estimate locally optimized model from inliers.
                if (support.num_inliers > Estimator::MinNumSamples() &&
                        support.num_inliers >= LocalEstimator::MinNumSamples())
                {
                    X_inlier.clear();
                    Y_inlier.clear();
                    X_inlier.reserve(support.num_inliers);
                    Y_inlier.reserve(support.num_inliers);
                    for (size_t i = 0; i < residuals.size(); ++i)
                    {
                        if (residuals[i] <= max_residual)
                        {
                            X_inlier.push_back(X[i]);
                            Y_inlier.push_back(Y[i]);
                        }
                    }

                    const std::vector<typename LocalEstimator::M_t> local_models =
                        local_estimator.Estimate(X_inlier, Y_inlier);

                    for (const auto& local_model : local_models)
                    {
                        local_estimator.Residuals(X, Y, local_model, &residuals);
                        CHECK_EQ(residuals.size(), X.size());

                        const auto local_support =
                            support_measurer.Evaluate(residuals, max_residual);

                        // Check if non-locally optimized model is better.
                        if (support_measurer.Compare(local_support, support))
                        {
                            best_support = local_support;
                            best_model = local_model;
                            best_model_is_local = true;
                        }
                    }
                }

                if (report.num_trials >= options_.min_num_trials)
                {
                    dyn_max_num_trials =
                        RANSAC<Estimator, SupportMeasurer, Sampler>::ComputeNumTrials(
                            best_support.num_inliers, num_samples, options_.confidence);
                }
            }

            if (report.num_trials >= dyn_max_num_trials)
            {
                abort = true;
                break;
            }
        }
    }

    report.support = best_support;
    report.model = best_model;

    // No valid model was found
    if (report.support.num_inliers < estimator.MinNumSamples())
    {
        return report;
    }

    report.success = true;

    // Determine inlier mask. Note that this calculates the residuals for the
    // best model twice, but saves to copy and fill the inlier mask for each
    // evaluated model. Some benchmarking revealed that this approach is faster.

    if (best_model_is_local)
    {
        local_estimator.Residuals(X, Y, report.model, &residuals);
        //if(options_.max_error == 8.0)
        //{
        //    LOG(INFO) << "Local Models is more accurate" << std::endl;
        //}
    }
    else
    {
        estimator.Residuals(X, Y, report.model, &residuals);
        //if(options_.max_error == 8.0)
        //{
        //    LOG(INFO) << "P3P is more accurate" << std::endl;
        //}
    }

    CHECK_EQ(residuals.size(), X.size());

    report.inlier_mask.resize(num_samples);
    for (size_t i = 0; i < residuals.size(); ++i)
    {
        if (residuals[i] <= max_residual)
        {
            report.inlier_mask[i] = true;
        }
        else
        {
            report.inlier_mask[i] = false;
        }
    }

    return report;
}

}  // namespace colmap

#endif  // COLMAP_SRC_OPTIM_LORANSAC_H_
